#这是原论文的版本
import torch
import torch.nn as nn
import torch.nn.functional as F
from networks.base_networks import Encoder_MDCBlock1, Decoder_MDCBlock1

def make_model(args, parent=False):
    return Net()
#密集连接的一个模块
class make_dense(nn.Module):
  def __init__(self, nChannels, growthRate, kernel_size=3):
    super(make_dense, self).__init__()
    #下面这个虽然用的是普通的卷积函数，但实际上是一个密集连接的卷积，具体不同体现在forward的实现中。growthRate在密集连接中称为增长率，也就是经过这个卷积层处理后
    #新的特征图通道数必须是输入通道数+增长率（增长率是整数，不是那种百分比）
    self.conv = nn.Conv2d(nChannels, growthRate, kernel_size=kernel_size, padding=(kernel_size-1)//2, bias=False)
  def forward(self, x):
      #这里的前向传播就是把处理过得到的growthRate个通道和输入的通道进行拼接
    out = F.relu(self.conv(x))
    out = torch.cat((x, out), 1)
    return out


class ConvLayer(nn.Module):
    #通过计算可得经过这一自定义卷积步长为1的卷积核得到的输入和输出的尺寸相同
    def __init__(self, in_channels, out_channels, kernel_size, stride):
        super(ConvLayer, self).__init__()
        reflection_padding = kernel_size // 2
        self.reflection_pad = nn.ReflectionPad2d(reflection_padding) #反射填充，可以使得之后卷积得到的图像在边缘上更加自然
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)

    def forward(self, x):
        out = self.reflection_pad(x)
        out = self.conv2d(out)
        return out

#这名字，上采样卷积层，大概就是反卷积的模块了
class UpsampleConvLayer(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride):
      super(UpsampleConvLayer, self).__init__()
      self.conv2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride)

    def forward(self, x):
        out = self.conv2d(x)
        return out

#这个是残差块
class ResidualBlock(torch.nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1) #他自定义的卷积层比提供的高级在有反射填充
        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.relu = nn.PReLU()  #用的PReLU
        #残差块的设计，输入为x，处理后结果为F（x），输出为x+F（x）
        #MSBDN对于残差块做出了一些改动，首先令F（x）成为两层处理的结果，我估计是可能层数太多了，然后令F（x）*0.1，这是为什么？
#我最能接受的说法就是使得残差对原始图像的影响变小，目的就是暗示图像微调即可
    def forward(self, x):
        residual = x #以自身也就是输入初始化残差
        out = self.relu(self.conv1(x))
        out = self.conv2(out) * 0.1
        out = torch.add(out, residual)
        return out
#这个大概是网络本体了
class Net(nn.Module):
    def __init__(self, res_blocks=18):
        super(Net, self).__init__()
        #步长为1自定义卷积处理
        self.conv_input = ConvLayer(3, 16, kernel_size=11, stride=1)
        #nn.Sequential可以构建一个简单的神经网络，如下的代码，dense0就是一个简单的网络结构，使用的时候直接dense0（x）
        #三次残差处理
        self.dense0 = nn.Sequential(
            ResidualBlock(16),
            ResidualBlock(16),
            ResidualBlock(16)
        )#这里已经6层卷积了
        #步长为2自定义卷积处理
        self.conv2x = ConvLayer(16, 32, kernel_size=3, stride=2)
        #编码器处理
        #输入通道数32，特征融合的层级是2，然后融合模式为模式2，后面的融合模式都是模式2
        self.fusion1 = Encoder_MDCBlock1(32, 2, mode='iter2')
        #三次残差块处理
        self.dense1 = nn.Sequential(
            ResidualBlock(32),
            ResidualBlock(32),
            ResidualBlock(32)
        )
        #步长为2卷积处理
        self.conv4x = ConvLayer(32, 64, kernel_size=3, stride=2)
        #编码器再处理
        self.fusion2 = Encoder_MDCBlock1(64, 3, mode='iter2')
        #三次残差块处理
        self.dense2 = nn.Sequential(
            ResidualBlock(64),
            ResidualBlock(64),
            ResidualBlock(64)
        )
        #步长为2卷积处理
        self.conv8x = ConvLayer(64, 128, kernel_size=3, stride=2)
        #编码器处理理
        self.fusion3 = Encoder_MDCBlock1(128, 4, mode='iter2')
        #依旧残差块
        self.dense3 = nn.Sequential(
            ResidualBlock(128),
            ResidualBlock(128),
            ResidualBlock(128)
        )
       #步长为2卷积处理
        self.conv16x = ConvLayer(128, 256, kernel_size=3, stride=2)
        #依旧编码器
        self.fusion4 = Encoder_MDCBlock1(256, 5, mode='iter2')
        #编码器部分做了四次实现一致的操作，就是残差块处理后跟上自定义卷积处理再跟上编码器的特征融合

        #这下面代码干嘛用的
        #self.dense4 = Dense_Block(256, 256)


        self.dehaze = nn.Sequential()#这代码构建了一个空的神经网络
        for i in range(0, res_blocks):
            self.dehaze.add_module('res%d' % i, ResidualBlock(256)) #这代码在为上面的神经网络加结构，就是猛猛加残差块，res_blocks这个是默认参数默认是18
        # 这下面应该是开始重构图像了
        #反卷积一次
        self.convd16x = UpsampleConvLayer(256, 128, kernel_size=3, stride=2)
        #依旧残差块
        self.dense_4 = nn.Sequential(
            ResidualBlock(128),
            ResidualBlock(128),
            ResidualBlock(128)
        )
        #解码器
        self.fusion_4 = Decoder_MDCBlock1(128, 2, mode='iter2')
        #上采样
        self.convd8x = UpsampleConvLayer(128, 64, kernel_size=3, stride=2)
        #残差块
        self.dense_3 = nn.Sequential(
            ResidualBlock(64),
            ResidualBlock(64),
            ResidualBlock(64)
        )
        #解码器
        self.fusion_3 = Decoder_MDCBlock1(64, 3, mode='iter2')
        #上采样
        self.convd4x = UpsampleConvLayer(64, 32, kernel_size=3, stride=2)
        #残差块
        self.dense_2 = nn.Sequential(
            ResidualBlock(32),
            ResidualBlock(32),
            ResidualBlock(32)
        )
        #解码器
        self.fusion_2 = Decoder_MDCBlock1(32, 4, mode='iter2')
        #上采样
        self.convd2x = UpsampleConvLayer(32, 16, kernel_size=3, stride=2)
        #残差块
        self.dense_1 = nn.Sequential(
            ResidualBlock(16),
            ResidualBlock(16),
            ResidualBlock(16)
        )
        #解码器
        self.fusion_1 = Decoder_MDCBlock1(16, 5, mode='iter2')
        #输出的通道为3，是因为输入为rgb图像，输入为3通道
        self.conv_output = ConvLayer(16, 3, kernel_size=3, stride=1)


    def forward(self, x):
        res1x = self.conv_input(x) #这里用res1x作为起始的输入是有道理的，因为这个res1x是可以被梯度调节从而达到残差的一个选择作用
        feature_mem = [res1x]
        #这下面残差连接后又跟一个res1x，其实是构建了一个更大尺度上的残差块
        #这里是双重残差连接的效果，最大的作用就是对抗特征稀释，图像复原任务很重要的一点就是原始特征必须得到充分的保留
        x = self.dense0(res1x) + res1x
        res2x = self.conv2x(x)
        res2x = self.fusion1(res2x, feature_mem) #这里还强化了特征表达，吸收了升维丢失的信息
        feature_mem.append(res2x)
        #来解说一些MSBDN中编码器的这个基本模块，这个双重残差操作不改变图像通道数
        res2x =self.dense1(res2x) + res2x #这样的双重残差连接可以对抗特征的稀释，单纯经过残差块处理原始特征已经被稀释了多次，但是加上的res2x只被稀释了两次
        #这次的卷积是步长为2的卷积会令通道数翻倍，翻倍通道数最简单的理解就是获得更多更多样性的特征
        res4x =self.conv4x(res2x)
        #这里又开始补偿丢失的信息了
        res4x = self.fusion2(res4x, feature_mem)
        feature_mem.append(res4x)
        res4x = self.dense2(res4x) + res4x

        res8x = self.conv8x(res4x)
        res8x = self.fusion3(res8x, feature_mem)
        feature_mem.append(res8x)
        res8x = self.dense3(res8x) + res8x  #这里是一个残差连接同时也为了增强特征提取能力，让这个尺度的特征猛猛提取

        res16x = self.conv16x(res8x) #升维了
        res16x = self.fusion4(res16x, feature_mem) #融合特征
        #res16x = self.dense4(res16x)
        #这是一个解码器和编码器的桥梁，编码器输出的高层语义可以说抽象极了，直接输入给解码器，解码器可能读不懂，所以经过多层残差块进行一个特征的整理和抽象
        res_dehaze = res16x
        #这里的缩放系数完全可以替换成一个可学习的参数
        in_ft = res16x*2  #这里乘以2是因为深层语义通常值较小，乘2可以起到放大效果，虽然看起来没大多少，但是可能一个乘2就让很多值可以跨过激活函数，梯度层面产生一个累积
        res16x = self.dehaze(in_ft) + in_ft - res_dehaze #这里从代码角度上来说确实多余了，但是作者方便理解才这么写的
        feature_mem_up = [res16x]

        res16x = self.convd16x(res16x) #这是一个转置卷积优点是具有可以学习的参数
        #这里方法的意思是对res16x上采样，采样到res8x的尺寸
        #这里进行两次不同的上采样的原因是单纯转置卷积处理可能得到的图像尺寸会有误差，所以又加入了人工的双线性插值来处理使得尺寸更精确
        res16x = F.upsample(res16x, res8x.size()[2:], mode='bilinear')
        #简单相加为什么不会破坏原始有的特征
        #很简单我已经逐渐理解一切了，两个相加的特征图可以通过反向传播追溯梯度，达到调整两个特征图的输出，实现一个特征的选择，反向传播很神奇吧。
        res8x = torch.add(res16x, res8x) #融合特征，融合特征有很多好的办法，最简单的就是相加，网络之后会自己识别特征
        #卷积不可逆，单纯的上采样无法还原浅层丢失的信息，使用残差块提升特征表达能力然后加上原来的浅层信息，达到深层信息与浅层信息的融合
        res8x = self.dense_4(res8x) + res8x - res16x  #后续的res8x（新）-res16x其实就是res8x未融合的时候
        res8x = self.fusion_4(res8x, feature_mem_up)  #进一步融合浅层和深层细节
        feature_mem_up.append(res8x)

        res8x = self.convd8x(res8x)
        res8x = F.upsample(res8x, res4x.size()[2:], mode='bilinear')
        res4x = torch.add(res8x, res4x)
        res4x = self.dense_3(res4x) + res4x - res8x
        res4x = self.fusion_3(res4x, feature_mem_up)
        feature_mem_up.append(res4x)

        res4x = self.convd4x(res4x)
        res4x = F.upsample(res4x, res2x.size()[2:], mode='bilinear')
        res2x = torch.add(res4x, res2x)
        res2x = self.dense_2(res2x) + res2x - res4x 
        res2x = self.fusion_2(res2x, feature_mem_up)
        feature_mem_up.append(res2x)

        res2x = self.convd2x(res2x)
        res2x = F.upsample(res2x, x.size()[2:], mode='bilinear')
        x = torch.add(res2x, x)
        x = self.dense_1(x) + x - res2x 
        x = self.fusion_1(x, feature_mem_up)

        x = self.conv_output(x)

        return x
